{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from langchain.chains.combine_documents import create_stuff_documents_chain\n",
    "from langchain_community.embeddings import SparkLLMTextEmbeddings\n",
    "from langchain_community.llms.sparkllm import SparkLLM\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_core.messages import HumanMessage\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "\n",
    "os.environ[\"IFLYTEK_SPARK_APP_ID\"] = \"849d9562\"\n",
    "os.environ[\"IFLYTEK_SPARK_API_KEY\"] = \"e33637f01cebee4a4675eb479d77df12\"\n",
    "os.environ[\"IFLYTEK_SPARK_API_SECRET\"] = \"ZmI0NGE5ZGUwNzQwMGZjYWFiYjg4MTIx\"\n",
    "xinghuo_llm = SparkLLM(temperature=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "SYSTEM_TEMPLATE = \"\"\"\n",
    "根据以下上下文回答用户的问题。\n",
    "如果上下文不包含与问题相关的任何信息，请不要编造一些东西，只说“我不知道”：\n",
    "\n",
    "<context>\n",
    "{context}\n",
    "</context>\n",
    "\"\"\"\n",
    "\n",
    "question_answering_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            SYSTEM_TEMPLATE,\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"messages\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "document_chain = create_stuff_documents_chain(xinghuo_llm, question_answering_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatPromptValue(messages=[SystemMessage(content='\\n根据以下上下文回答用户的问题。\\n如果上下文不包含与问题相关的任何信息，请不要编造一些东西，只说“我不知道”：\\n\\n<context>\\n1\\n</context>\\n'), HumanMessage(content='111')])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "question_answering_prompt.format_prompt(context=\"1\",messages=[({\"role\": \"user\", \"content\": 111})])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "embeddings = SparkLLMTextEmbeddings(\n",
    "    spark_app_id=\"849d9562\",\n",
    "    spark_api_key=\"e33637f01cebee4a4675eb479d77df12\",\n",
    "    spark_api_secret=\"ZmI0NGE5ZGUwNzQwMGZjYWFiYjg4MTIx\",\n",
    ")\n",
    "\n",
    "vectordb_load = Chroma(\n",
    "    persist_directory=\"./\",\n",
    "    embedding_function=embeddings\n",
    ")\n",
    "retriever = vectordb_load.as_retriever()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VectorStoreRetriever(tags=['Chroma', 'SparkLLMTextEmbeddings'], vectorstore=<langchain_community.vectorstores.chroma.Chroma object at 0x000001FA3F385C90>)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docs = retriever.invoke(\"萧炎\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'AI: 当然可以，请提供更多信息。'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "document_chain.invoke(\n",
    "    {\n",
    "        \"context\": docs,\n",
    "        \"messages\": [\n",
    "            HumanMessage(content=\"LangSmith 可以帮助测试我的 LLM 应用程序吗?\")\n",
    "        ],\n",
    "    }\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
